'''
@Descripttion: 
@Author: Chen Chenxi
@Date: 2019-11-22 11:18:27
@LastEditTime: 2019-11-22 12:57:28
'''
from spatial_cnn import Spatial_CNN
import torch.tensor as tensor
from extract_frames import get_frames
import os
import cv2
from network import *
import torch
import torchvision.transforms as transforms
from PIL import Image
import torchvision
import numpy as np

#load spatial ConvNet model(trained)
model_path = "E:/ResSpatial.tar"
checkpoint = torch.load(model_path,map_location=torch.device('cpu'))
spatial_model = resnet101(pretrained= True, channel=3)
spatial_model.load_state_dict(checkpoint['state_dict'])
spatial_model.eval()

def input_video_convert(video_path):
    SAMPLE_NUM_PRE_SECOND = 2
    img_dir = os.path.dirname(video_path) + "/sample_img/"
    delete_old_img(img_dir)
    if not os.path.exists(img_dir):
        os.makedirs(img_dir)
    get_frames(video_path, SAMPLE_NUM_PRE_SECOND, img_dir)
    return img_dir


def load_label_list():
    with open("../UCF_list/classInd.txt", 'r') as f:
        label_list = [a[:-1] for a in f]
    return label_list

def predictor(video_path):
    img_dir = input_video_convert(video_path)

    frame_transform = transforms.Compose([
                    transforms.Resize([224,224]),
                    transforms.ToTensor(),
                    transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.225])
                    ])
    
    # #load spatial ConvNet model(trained)
    # checkpoint = torch.load("../record/spatial/model_best.pth.tar",map_location=torch.device('cpu'))
    # spatial_model = resnet101(pretrained= True, channel=3)
    # spatial_model.load_state_dict(checkpoint['state_dict'])
    # spatial_model.eval()


    img_l = os.listdir(img_dir)
    fram_nums = len(img_l)
    # spatial model predict
    with torch.no_grad():
        label_list = load_label_list()
        output = np.zeros((1,101))
        for i in range(1,fram_nums):
            image_data = Image.open(img_dir + "/test_"+str(i) + ".jpg")
            image_data = frame_transform(image_data)
            image_data = image_data.unsqueeze(0)
            output += spatial_model(image_data).data.cpu().numpy()
            # print(spatial_model(image_data).data.cpu().argmax(dim=1))
            # print(output.argmax(dim=1))
            # print(output.size())
        return label_list[np.argmax(output)]


def delete_old_img(img_dir):
    img_list = [img_dir + i for i in os.listdir(img_dir)]
    for i in img_list:
        os.remove(i)